package zdgm

import (
	"crypto/rand"
	"encoding/hex"
	"errors"
	"math/big"

	"github.com/tjfoc/gmsm/sm2"
	"github.com/tjfoc/gmsm/sm4"
)

type Sm4Mode int

const (
	SM4CBC Sm4Mode = iota
	SM4ECB
)

type ZdSm2 struct {
	privateKey *sm2.PrivateKey
}

type ZdSm4 struct {
	sm4Key      []byte
	encryptMode Sm4Mode
}

func NewZdSm2(publicXKey, publicYKey, privateKey string) (*ZdSm2, error) {
	privateKeyArray, err := hex.DecodeString(privateKey)
	publicXKeyArray, err := hex.DecodeString(publicXKey)
	publicYKeyArray, err := hex.DecodeString(publicYKey)

	if err != nil {
		return nil, err
	}

	c := sm2.P256Sm2()
	priv := new(sm2.PrivateKey)
	priv.PublicKey.Curve = c
	priv.D = new(big.Int).SetBytes(privateKeyArray)
	priv.PublicKey.X = new(big.Int).SetBytes(publicXKeyArray)
	priv.PublicKey.Y = new(big.Int).SetBytes(publicYKeyArray)
	zdSm2 := new(ZdSm2)
	zdSm2.privateKey = priv
	return zdSm2, nil
}

func NewZdSm4(sm4Key string, mode Sm4Mode) (*ZdSm4, error) {
	sm4KeyArray, err := hex.DecodeString(sm4Key)
	if err != nil {
		return nil, err
	}

	if mode != SM4CBC && mode != SM4ECB {
		return nil, errors.New("unsupported sm4 mode")
	}

	sm4 := &ZdSm4{sm4Key: sm4KeyArray, encryptMode: mode}
	return sm4, nil
}

func (zdSm2 *ZdSm2) Sm2Encrypt(originData *[]byte) ([]byte, error) {
	pub := zdSm2.privateKey
	d0, err := pub.EncryptAsn1(*originData, rand.Reader)
	if err != nil {
		return nil, err
	}
	return d0, nil
}

func (zdSm2 *ZdSm2) Sm2Decrypt(encryptData *[]byte) ([]byte, error) {
	data, err := zdSm2.privateKey.DecryptAsn1(*encryptData)
	if err != nil {
		return nil, err
	}
	return data, nil
}

func (zdSm4 *ZdSm4) Sm4Encrypt(originData *[]byte) ([]byte, error) {

	var data []byte
	var err error
	if zdSm4.encryptMode == SM4ECB {
		data, err = sm4.Sm4Ecb(zdSm4.sm4Key, *originData, true)
	} else {
		data, err = sm4.Sm4Cbc(zdSm4.sm4Key, *originData, true)
	}

	if err != nil {
		return nil, err
	}
	return data, nil
}

func (zdSm4 *ZdSm4) Sm4Decrypt(encryptData *[]byte) ([]byte, error) {
	var data []byte
	var err error
	if zdSm4.encryptMode == SM4ECB {
		data, err = sm4.Sm4Ecb(zdSm4.sm4Key, *encryptData, false)
	} else {
		data, err = sm4.Sm4Cbc(zdSm4.sm4Key, *encryptData, false)
	}
	if err != nil {
		return nil, err
	}
	return data, nil
}
